#include "expression.hpp"
#include "io.hpp"
#include "util.hpp"
#include "option_parse.hpp"
#include "perm_parser.hpp"
#include "call_rates.hpp"


void do_expression(int argc, char **argv)
{
    prog_options options = parse_expression_options(argc, argv);

    mapping_type t = get_mapping_type(options.map_file.c_str());
    if(t == COLOR)
        do_color_expression(options);
    else
        do_seq_expression(options);
}

void do_color_expression(prog_options &options)
{

    std::map<std::string, unsigned int> genome_size = build_fasta_index(options.fasta_file.c_str());

    unsigned int read_length = get_read_length(options.map_file.c_str());


    map_str_vafl_t expression;

    map_str_ui_t::iterator iter;

    for(iter = genome_size.begin(); iter != genome_size.end(); iter++)
        expression[iter->first].resize(iter->second);


    // load call rates data.
    array_type5 call_rates;
    array_type5 call_counts;
    boost::multi_array<int, 2> qual_bins;

    if(options.call_file != "")
    {
        read_call_counts_file(options.call_file.c_str(), call_counts);
        // TODO: UGLY: Is options.old_bin_cnt correctly set here?
        process_color_counts_to_rates(call_counts, call_rates, options.old_bin_cnt, options.new_bin_cnt, read_length, qual_bins);
    
        ColorRead read;

        ColorMappingParser parser(options.map_file.c_str());
        while( parser.next_read(read) )
        {
            if(options.expr_unique && read.num > 1)
                continue;

            // translate the quality score into our new score.
            for(unsigned int i = 0; i < read_length; i++)
                read.qual[i] = qual_bins[i][read.qual[i]];

            std::vector<double> align_scores(read.num);
            std::string dna_ref;
            
            align_scores.resize(read.num);

            double sum_scores  = 0;

            // calculate all alignment scores to mapping locations using genome seq.
            for(unsigned int j = 0; j < read.num; j++)
            {
                MapLoci &loc = read.map_locs[j];

                dna_ref = color_to_seq(loc.seq);

                double align_score = 1.0;

                for(unsigned int i = 0; i < dna_ref.size() - 1; i++)
                    align_score *= call_rates[i+1][read.qual[i+1]][cswap[dna_ref[i]]][cswap[dna_ref[i+1]]][iswap[read.seq[i+1]]];

                if( align_score > 1 || isnan(align_score) )
                    error("align_score > 1 or align_score is NaN");
                
                align_scores[j] = align_score;
                sum_scores  += align_score;
            }

            if(options.expr_random)
            {
                // choose a random site based on alignment score.
                double R = ((double)rand()) / RAND_MAX * sum_scores;
                double sum = 0.0;
                unsigned int j;
                for(j = 0; j < read.num; j++)
                {
                    sum += align_scores[j];
                    if( sum > R )
                        break;
                }

                // we have chosen mapping j.
                MapLoci &loc = read.map_locs[j];
                for(unsigned int i = 0; i < read_length; i++)
                    expression[loc.id][loc.pos + i] += 1.0;
            }
            else if (options.expr_all)
            {
                // add a small amount to each proportional to mapping score.
                unsigned int j;
                for(j = 0; j < read.num; j++)
                {
                    MapLoci &loc = read.map_locs[j];
                    double delta = align_scores[j]/sum_scores;
                    for(unsigned int i = 0; i < read_length; i++)
                        expression[loc.id][loc.pos + i] += delta;
                }
            }
            else if (options.expr_unique)
            {
                MapLoci &loc = read.map_locs[0];
                for(unsigned int i = 0; i < read_length; i++)
                    expression[loc.id][loc.pos + i] += 1.0;
            }
        }
    }
    else
    {
        ColorRead read;

        ColorMappingParser parser(options.map_file.c_str());
        while( parser.next_read(read) )
        {
            if(options.expr_unique && read.num > 1)
                continue;

            std::vector<double> align_scores;
            std::string dna_ref;
            
            align_scores.resize(read.num);
            dna_ref.resize(read.num);

            //double sum_scores  = 2;

            if(options.expr_random)
            {
                // choose a random site based on alignment score.
                unsigned int j = (unsigned int)floor( (rand() / (double)RAND_MAX) * read.map_locs.size());

                // we have chosen mapping j.
                MapLoci &loc = read.map_locs[j];
                for(unsigned int i = 0; i < read_length; i++)
                    expression[loc.id][loc.pos + i] += 1.0;
            }
            else if (options.expr_all)
            {
                // add a small amount to each proportional to number of mapping locations.
                unsigned int j;
                for(j = 0; j < read.num; j++)
                {
                    MapLoci &loc = read.map_locs[j];
                    double delta = 1.0 / read.map_locs.size();
                    for(unsigned int i = 0; i < read_length; i++)
                        expression[loc.id][loc.pos + i] += delta;
                }
            }
            else if (options.expr_unique)
            {
                MapLoci &loc = read.map_locs[0];
                for(unsigned int i = 0; i < read_length; i++)
                    expression[loc.id][loc.pos + i] += 1.0;
            }
        }
    }

    write_expression(std::cout, expression);
}


void do_seq_expression(prog_options &options) 
{
    error("sequence expression not yet implemented.");
}
